﻿#include "ScaledShapeBasedMatching.h"
#include <QDir>
#include <QDebug>

namespace Qly {
ScaledShapeBasedMatching::ScaledShapeBasedMatching()
{
}

bool ScaledShapeBasedMatching::trainModel(cv::Mat modelImg, float min_scale, float max_scale, int num_scale, int num_features)
{
    if(modelImg.empty() || min_scale <= 0 || max_scale <= min_scale || num_features < 4)
    {
        return false;
    }
    preprocessModelImage(modelImg, 0, 0);
    shape_based_matching::shapeInfo_producer shapes(templateImg);
    shapes.produce_infos_logscales(min_scale, max_scale, num_scale);
    infos_have_templ.clear();
    //cv::Mat mask = cv::Mat(templateImg.size(), CV_8UC1, {255});
    if(pDetector)
    {
        delete pDetector;
    }
    pDetector = new line2Dup::Detector(num_features, {4, 8});
    for(auto& info: shapes.infos)
    {
        // template img, id, mask,
        //feature numbers(missing it means using the detector initial num)

        std::cout << "info.angle = "  << info.angle << ", info.scale = " << info.scale << std::endl;
        info.templ_id = pDetector->addTemplate(shapes.src_of(info),
                                              class_id,
                                              shapes.mask_of(info),
                                              int(num_features * info.scale));
        std::cout << "templ_id: " << info.templ_id << std::endl;

        // may fail when asking for too many feature_nums for small training img
        if(info.templ_id != -1)
        {  // only record info when we successfully add template
            infos_have_templ.push_back(info);
        }
    }
    if(infos_have_templ.empty()) return false;
    return true;
}

bool ScaledShapeBasedMatching::saveModel(std::string path, std::string modelName)
{
    if(pDetector == nullptr) return false;
    std::string fullModelName = path + "/" + modelName + ".yaml";
    std::string fullInfoName = path + "/" + modelName + "_info.yaml";

    QDir dir;
    dir.remove(QString::fromStdString(fullModelName));
    dir.remove(QString::fromStdString(fullInfoName));

    // save templates
    pDetector->writeClasses(fullModelName);
    shape_based_matching::shapeInfo_producer::save_infos(infos_have_templ, fullInfoName);

    std::cout << "save model end" << std::endl << std::endl;
    return true;
}

bool ScaledShapeBasedMatching::loadModel(std::string path, std::string modelName)
{
    //模板文件路径获取
    std::string fullModelName = path + "/" + modelName + ".yaml";
    std::string fullInfoName = path + "/" + modelName + "_info.yaml";

    class_ids.clear();
    class_ids.push_back("test");
    pDetector->readClasses(class_ids, fullModelName);
    infos_have_templ = shape_based_matching::shapeInfo_producer::load_infos(fullInfoName);
    if(pDetector->numTemplates() != 0)
    {
        return true;
    }
    return false;
}

std::vector<ShapeBasedMatching::MatchResult> ScaledShapeBasedMatching::runMatch(const cv::Mat &img, float score, size_t maxCount)
{
    std::vector<ScaledShapeBasedMatching::MatchResult> ret;
    if(img.empty() || !pDetector) return ret;
    cv::Mat mat = preprocessTestImage(img);
    //cv::imshow("111", mat);
    std::vector<line2Dup::Match> matches = pDetector->match(mat, score, class_ids);// 这里得到的 matches 已经是排好序的了

    qDebug() <<"pDetector->match size = " << matches.size();
    for(size_t i = 0; i < matches.size(); i++)
    {
        qDebug() << "scale = " << scale(matches[i]) << ", score = " << matches[i].similarity << ", id = " << matches[i].template_id;
    }
    refine(matches, maxCount, ret);
    return ret;
}

//void ScaledShapeBasedMatching::refine(const std::vector<line2Dup::Match> &matches, size_t maxCount, std::vector<MatchResult> &ret)
//{
//    ret.clear();
//    //将匹配到的最小外接矩形装入vector
//    std::vector<cv::Rect> vRects;
//    std::vector<float> vScores;
//    std::vector<float> vScale;
//    std::vector<bool> vDeleteIndexs;
//    for (size_t i = 0; i < matches.size() ; i++)
//    {
//        auto match = matches[i];
//        float sc = scale(match);
//        cv::Rect rect = boundRect(match, sc);
//        vRects.push_back(rect);
//        vScale.push_back(sc);
//        vScores.push_back(match.similarity);
//        vDeleteIndexs.push_back(false);
//    }
//    //剔除重复最小外接矩形
//    cv::Rect rect1, rect2;
//    int countDelete = 0;
//    for (size_t i = 0; i < vRects.size() - 1; i++)
//    {
//        if (vDeleteIndexs[i])
//            continue;
//        for (size_t j = i + 1; j < vRects.size(); j++)
//        {
//            if (vDeleteIndexs[j])
//                continue;
//            rect1 = vRects[i];
//            rect2 = vRects[j];
//            cv::Rect un = rect1 & rect2;
//            if ( un.empty() )//無交集
//                continue;
//            else if (rect1 == (rect1 & rect2) || rect2 == (rect1 & rect2)) //一個矩形包覆另一個
//            {
//                int iDeleteIndex =  (vScores[i] <= vScores[j] ? i : j); //匹配分数低的被剔除
//                vDeleteIndexs[iDeleteIndex] = true;
//                countDelete ++;
//            }
//            else//交點 > 0
//            {
//                double dArea = (rect1 & rect2).area();
//                if (dArea >= 0.5 * rect1.area() || dArea >= 0.5 * rect2.area())
//                {
//                    int iDeleteIndex = (vScores[i] <= vScores[j] ? i : j); //匹配分数低的被剔除
//                    vDeleteIndexs[iDeleteIndex] = true;
//                    countDelete ++;
//                }
//            }
//        }
//    }
//    int rectCount = vRects.size() - countDelete;
//    maxCount = std::min((int)maxCount, rectCount);

//    for(size_t i = 0; i < vRects.size(); i++)
//    {
//        if( !vDeleteIndexs[i] )
//        {
//            MatchResult mr;
//            mr.scale = vScale[i];
//            mr.rect = vRects[i];
//            mr.score = vScores[i];
//            ret.push_back(mr);
//        }
//    }
//}


cv::Rect ScaledShapeBasedMatching::boundRect(const line2Dup::Match &match, float scale)
{
    std::vector<line2Dup::Template> templ = pDetector->getTemplates(class_id, match.template_id);

    float x = match.x - templ[0].tl_x - 32;
    float y = match.y - templ[0].tl_y - 32;

    return cv::Rect(x, y, templateImg.cols * scale, templateImg.rows * scale);
}

float ScaledShapeBasedMatching::scale(const line2Dup::Match &match)
{
//    float scale = 1;
//    for(size_t i = 0; i < infos_have_templ.size(); i++)
//    {
//        if( infos_have_templ[i].templ_id == match.template_id )
//        {
//            scale = infos_have_templ[i].scale;
//            return scale;
//        }
//    }
//    return -1.0;
    return infos_have_templ[match.template_id].scale;
}

}

